#include "code.h"
#include "midcode.h"
#include "table.h"
#include "util.h"
#include "global.h"

// Basic block.
int bblock[MAX_FUNC_NUM][BBLOCK_NUM][2] = {0};
const char REGISTER_STRING[][4] = {
    "EAX", "ECX", "EDX", "EBX", "ESI", "EDI", "EBP",
};
RegVar reg_pool[REG_NUM] = {0};
// Number of being referenced.
static int ref_cnt[MAX_TOKEN_NUM] = {0};
// For the index of running stack and parameter.
static int running_stack_ref[MAX_TOKEN_NUM] = {0};
// Record the number of variables in each function.
static int func_var_num[MAX_FUNC_NUM] = {0};
// The first normal element's address.
static NormalElement *fne = &normal_table[0];


void print_line(char *s)
{
    fprintf(asmout, "%s\n", s);
}

static Register inRegister(NormalElement *ne) {
    int i;
    for (i = 0; i < REG_NUM; i++) {
        if (reg_pool[i].ne == ne) {
            return reg_pool[i].rg;
        }
    }
    return -1;
}

static Register inGlobalRegister(NormalElement *ne) {
    int i;
    for (i = EBX; i <= EDI; i++) {
        if (reg_pool[i].ne == ne) {
            return reg_pool[i].rg;
        }
    }
    return -1;
}

static char *toString(NormalElement *ne, char st[])
{
    int i, j, k;

    if (ne->clss != GLOBAL_VAR) {
        i = running_stack_ref[ne - fne];
        j = ne->scope->param_num;
        k = running_stack_ref[ne - fne];
        switch (ne->clss) {
            case PARAMETER:
                sprintf(st, "DWORD PTR [EBP + 0%xH]", (j - i + 1) << 2);
                break;
            case LOCAL_VAR:
            case TEMP_VAR:
                sprintf(st, "DWORD PTR [EBP - 0%xH]", (k + 1) << 2);
                break;
            case GLOBAL_CONST:
            case LOCAL_CONST:
                sprintf(st, "%d", ne->value.int_val);
                break;
            default:
                break;
        }
    }
    else {
        sprintf(st, "_%s", ne->name);
    }
    return st;
}
static char *operandToString(Operand or, int ln, char os[])
{
    int i, j, k;
    Register tr;
    switch (or.operandType) {
        case NUM_MB:
            sprintf(os, "0%xH", or.operandCont.constOp);
            break;
        case STR_MB:
            sprintf(os, "OFFSET c_str_%d", ln);
            break;
        case FUNC_MB:
            sprintf(os, "%s", or.operandCont.funcOp->name);
            break;
        case JUMP_MB:
            sprintf(os, "TARGET_%d", or.operandCont.midcOp);
            break;
        case GLOBAL_MB:
            tr = inRegister(or.operandCont.varOp);
            if (tr == -1) {
                sprintf(os, "_%s", or.operandCont.varOp->name);
            }
            else {
                sprintf(os, "%s", REGISTER_STRING[tr]);
            }
            break;
        case TMPVAR_MB:
            tr = inRegister(or.operandCont.varOp);
            if (tr == -1) {
                k = running_stack_ref[or.operandCont.varOp - fne];
                sprintf(os, "DWORD PTR [EBP - 0%xH]", (k + 1) << 2);
            }
            else {
                sprintf(os, "%s", REGISTER_STRING[tr]);
            }
            break;
        case LVAR_MB:
            tr = inRegister(or.operandCont.varOp);
            if (tr != -1) {
                sprintf(os, "%s", REGISTER_STRING[tr]);
            }
            else if (or.operandCont.varOp->clss == PARAMETER) {
                i = running_stack_ref[or.operandCont.varOp - fne];
                j = or.operandCont.varOp->scope->param_num;
                sprintf(os, "DWORD PTR [EBP + 0%xH]", (j - i + 1) << 2);
            }
            else {
                k = running_stack_ref[or.operandCont.varOp - fne];
                sprintf(os, "DWORD PTR [EBP - 0%xH]", (k + 1) << 2);
            }
            break;
        default:
            break;
    }
    return os;
}

void gen_bblock()
{
    int i, j, k;
    char mark[MIDCODE_NUM] = {0};
    for (i = 0; i < func_table_p; i++) {
        for (j = func_mid_map[i][0] + 1, mark[j] = 1; j <= func_mid_map[i][1]; j++) {
            if (isJumpOp(midcode_table[j].op0) == TRUE) {
                mark[j + 1] = 1;
                if (midcode_table[j].op0 == JMP_OP) {
                    mark[midcode_table[j].op1.operandCont.midcOp] = 1;
                }
                else {
                    mark[midcode_table[j].op3.operandCont.midcOp] = 1;
                }
            }
        }
        k = 0;
        for (j = func_mid_map[i][0] + 1; j <= func_mid_map[i][1]; j++) {
            if (mark[j] == 1) {
                bblock[i][k][0] = j++;
                while (mark[j] == 0 && j <= func_mid_map[i][1]) {
                    j++;
                }
                j--;
                bblock[i][k][1] = j;
                k++;
            }
        }
    }
}


void init_reg_pool()
{
    int i;
    for (i = 0; i < REG_NUM; i++) {
        reg_pool[i].rg = i;
        reg_pool[i].rs = FREE;
        reg_pool[i].ne = NULL;
    }
}

static void inc_ref(Operand or)
{
    if ((or.operandType == LVAR_MB || or.operandType == TMPVAR_MB) &&
        (or.operandCont.varOp->clss != PARAMETER)) {
            ref_cnt[(or.operandCont.varOp - fne)]++;
    }
}
static void release_reg(int index)
{
    char t_str[MAX_TOKEN_LEN] = {0};
    if (reg_pool[index].rs == OCCUPIED) {
        if (reg_pool[index].ne != NULL) {
            fprintf(asmout, "MOV\t%s, %s\n", toString(reg_pool[index].ne, t_str), REGISTER_STRING[reg_pool[index].rg]);
        }
        reg_pool[index].ne = NULL;
        reg_pool[index].rs = FREE;
    }
}
static void release_tmp_reg()
{
    int i;

    for (i = EAX; i <= EDX; i++) {
        release_reg(i);
    }
}
// @param: block, func_mid_map
void global_reg_dist(int block[/*2*/])
{
    int i, j, k;
    int ref_num[3] = {0};                        // Biggest, medium, smallest
    NormalElement *ne[3] = {0};
    FuncElement *t_fe = midcode_table[block[0]].op1.operandCont.funcOp;
    for (i = block[0] + 1; i <= block[1]; i++) {
        inc_ref(midcode_table[i].op1);
        inc_ref(midcode_table[i].op2);
        inc_ref(midcode_table[i].op3);
    }
    // Get the first three normal_elements that were referenced most.
    for (i = 0; i < normal_table_p; i++) {
        if ((fne + i)->scope == t_fe) {
            for (j = 0; j < 3; j++) {
                if (ref_cnt[i] > ref_num[j]) {
                    for (k = 2; k > j; k--) {
                        ref_num[k] = ref_num[k - 1];
                        ne[k] = ne[k - 1];
                    }
                    ref_num[j] = ref_cnt[i];
                    ne[j] = fne + i;
                    break;
                }
            }
        }
    }
    for (j = EBX; j <= EDI && ne[j - EBX] != NULL; j++) {
        reg_pool[j].ne = ne[j - EBX];
        reg_pool[j].rs = OCCUPIED;
        reg_pool[j].rg = j;
    }
}

Register tmp_reg_dist(NormalElement *ne, Bool assignee)
{
    int i;
    Bool gotcha = FALSE;
    char t_str[MAX_TOKEN_LEN] = {0};
    int j = -1;
    static int index = EAX;
    for (i = EBX; i <= EDI; i++) {
        if (reg_pool[i].ne == ne) {
            return reg_pool[i].rg;
        }
    }
    for (i = EAX; i <= EDX; i++) {
        // "ne" is already in register.
        if (reg_pool[i].ne == ne) {
            return reg_pool[i].rg;
        }
        if (reg_pool[i].rs == FREE && gotcha == FALSE) {
            gotcha = TRUE;
            j = i;
        }
    }
    // Found free register.
    if (gotcha == TRUE) {
        reg_pool[j].ne = ne;
        reg_pool[j].rs = OCCUPIED;
        if (assignee == FALSE) {
            fprintf(asmout, "MOV\t%s, %s\n", REGISTER_STRING[reg_pool[j].rg], toString(ne, t_str));
        }
        return reg_pool[j].rg;
    }
    else {
        release_reg(index);
        reg_pool[index].ne = ne;
        reg_pool[index].rs = OCCUPIED;
        j = index;
        if (index != EDX) {
            index++;
        }
        else {
            index = EAX;
        }
        return reg_pool[j].rg;
    }
}


/////////////////////////////
// Generate assembly code. //
/////////////////////////////

// @param: func, func_mid_map
static void gen_stack_ref()
{
    int i;
    int j = 0;
	int k = 0;                     // Variable position.
    int p = 0;                     // Parameter position.
    FuncElement *current_func = normal_table[0].scope;
    if (current_func != NULL) {
        func_var_num[0] = 0;
        k = 1;
    }
    for (i = 0; i < normal_table_p; i++) {
        j = 0;
        p = 0;
        while (i < normal_table_p && normal_table[i].scope == current_func) {
            if (normal_table[i].clss == LOCAL_VAR || normal_table[i].clss == TEMP_VAR) {
                running_stack_ref[i] = j;
                j++;
            }
            else if (normal_table[i].clss == PARAMETER) {
                running_stack_ref[i] = p;
                p++;
            }
            else {
                running_stack_ref[i] = -1;
            }
            i++;
        }
		func_var_num[k++] = j;
        current_func = normal_table[i--].scope;
    }
}

static void add_inst(int i)
{
    char t_str[MAX_TOKEN_LEN] = {0};
    Register tr = tmp_reg_dist(midcode_table[i].op3.operandCont.varOp, TRUE);
    fprintf(asmout, "MOV\t%s, %s\n", REGISTER_STRING[tr], operandToString(midcode_table[i].op1, i, t_str));
    fprintf(asmout, "ADD\t%s, %s\n", REGISTER_STRING[tr], operandToString(midcode_table[i].op2, i, t_str));
}
static void sub_inst(int i)
{
    char t_str[MAX_TOKEN_LEN] = {0};
    Register tr = tmp_reg_dist(midcode_table[i].op3.operandCont.varOp, TRUE);
    fprintf(asmout, "MOV\t%s, %s\n", REGISTER_STRING[tr], operandToString(midcode_table[i].op1, i, t_str));
    fprintf(asmout, "SUB\t%s, %s\n", REGISTER_STRING[tr], operandToString(midcode_table[i].op2, i, t_str));
}
static void mul_inst(int i)
{
    // IMUL SRC, SRC can't be imm type
    // EAX and ECX don't connect to any normal element.
    // EDX is going to be abandoned...
    char t_str[MAX_TOKEN_LEN] = {0};
    Register op3_rg;
    release_reg(EAX);
    release_reg(ECX);
    release_reg(EDX);

    reg_pool[EAX].ne = midcode_table[i].op3.operandCont.varOp;
    reg_pool[EAX].rs = OCCUPIED;
    fprintf(asmout, "MOV\t%s, %s\n", REGISTER_STRING[EAX], operandToString(midcode_table[i].op1, i, t_str));
    fprintf(asmout, "MOV\t%s, %s\n", REGISTER_STRING[ECX], operandToString(midcode_table[i].op2, i, t_str));
    fprintf(asmout, "XOR\t%s, %s\n", REGISTER_STRING[EDX], REGISTER_STRING[EDX]);
    fprintf(asmout, "IMUL\t%s\n", REGISTER_STRING[ECX]);
    op3_rg = inGlobalRegister(reg_pool[EAX].ne);
    if (op3_rg != -1) {
        fprintf(asmout, "MOV\t%s, %s\n", REGISTER_STRING[op3_rg], REGISTER_STRING[EAX]);
        release_reg(EAX);
    }
}
static void div_inst(int i)
{
    char t_str[MAX_TOKEN_LEN] = {0};
    Register op3_rg;
    release_reg(EAX);
    release_reg(ECX);
    release_reg(EDX);

    reg_pool[EAX].ne = midcode_table[i].op3.operandCont.varOp;
    reg_pool[EAX].rs = OCCUPIED;
    fprintf(asmout, "MOV\t%s, %s\n", REGISTER_STRING[EAX], operandToString(midcode_table[i].op1, i, t_str));
    fprintf(asmout, "MOV\t%s, %s\n", REGISTER_STRING[ECX], operandToString(midcode_table[i].op2, i, t_str));
    fprintf(asmout, "XOR\t%s, %s\n", REGISTER_STRING[EDX], REGISTER_STRING[EDX]);
    fprintf(asmout, "IDIV\t%s\n", REGISTER_STRING[ECX]);
    op3_rg = inGlobalRegister(reg_pool[EAX].ne);
    if (op3_rg != -1) {
        fprintf(asmout, "MOV\t%s, %s\n", REGISTER_STRING[op3_rg], REGISTER_STRING[EAX]);
        release_reg(EAX);
    }
}
static void if_inst(int i, Opcode oc)
{
    char J_INST[][4] = {"JB", "JBE", "JA", "JAE", "JNE", "JE", };
    char t_str[MAX_TOKEN_LEN] = {0};
    release_reg(EAX);
    release_reg(ECX);

    fprintf(asmout, "MOV\t%s, %s\n", REGISTER_STRING[EAX], operandToString(midcode_table[i].op1, i, t_str));
    fprintf(asmout, "MOV\t%s, %s\n", REGISTER_STRING[ECX], operandToString(midcode_table[i].op2, i, t_str));
    fprintf(asmout, "CMP\t%s, %s\n", REGISTER_STRING[EAX], REGISTER_STRING[ECX]);
    release_tmp_reg();
    fprintf(asmout, "%s\t%s\n", J_INST[oc - IFLT_OP], operandToString(midcode_table[i].op3, i, t_str));
}
static void push_inst(int i)
{
    char t_str[MAX_TOKEN_LEN] = {0};
    fprintf(asmout, "PUSH\t%s\n", operandToString(midcode_table[i].op1, i, t_str));
}

static void call_inst(int i)
{
    //char t_str[MAX_TOKEN_LEN] = {0};

    release_tmp_reg();
    fprintf(asmout, "CALL\t_%s\n", midcode_table[i].op1.operandCont.funcOp->name);
    reg_pool[EAX].rs = OCCUPIED;
    reg_pool[EAX].ne = midcode_table[i].op2.operandCont.varOp;
}
static void jmp_inst(int i)
{
    char t_str[MAX_TOKEN_LEN] = {0};
    release_tmp_reg();
    fprintf(asmout, "JMP\t%s\n", operandToString(midcode_table[i].op1, i, t_str));
}
static void mov_inst(int i)
{
    char t_str[MAX_TOKEN_LEN] = {0};
    Register tr = tmp_reg_dist(midcode_table[i].op1.operandCont.varOp, TRUE);
    fprintf(asmout, "MOV\t%s, %s\n", REGISTER_STRING[tr], operandToString(midcode_table[i].op2, i, t_str));
}
static void rev_inst(int i)
{
    char t_str[MAX_TOKEN_LEN] = {0};
    release_reg(EAX);
    fprintf(asmout, "MOV\t%s, %s\n", REGISTER_STRING[EAX], operandToString(midcode_table[i].op1, i, t_str));
}
static void ret_inst(int i)
{
    fprintf(asmout, "JMP\t%s\n", "TARGET_END");
}
static void prf_inst(int i)
{
    char t_str[MAX_TOKEN_LEN] = {0};
    if (midcode_table[i].op1.operandType == STR_MB) {
        fprintf(asmout, "PUSH\tOFFSET c_str_%d\n", i);

        release_tmp_reg();
        fprintf(asmout, "CALL\tprintf\n");
        fprintf(asmout, "ADD\tESP, 04H\n");
    }
    else/* if (midcode_table[i].op1.operandType = NUM_MB) */{
        release_reg(EAX);
        fprintf(asmout, "MOV\t%s, %s\n", REGISTER_STRING[EAX], operandToString(midcode_table[i].op1, i, t_str));
        fprintf(asmout, "PUSH\t%s\n", REGISTER_STRING[EAX]);
        fprintf(asmout, "PUSH\tOFFSET %s\n", "c_str_int_out");

        release_tmp_reg();
        fprintf(asmout, "CALL\tprintf\n");
        fprintf(asmout, "ADD\tESP, 0%xH\n", 8);
    }
}
static void scf_inst(int i)
{

    char t_str[MAX_TOKEN_LEN] = {0};
    Register tr = inRegister(midcode_table[i].op1.operandCont.varOp);
    if (tr != -1) {
        release_reg(tr);
    }
    fprintf(asmout, "LEA\tEAX, %s\n", operandToString(midcode_table[i].op1, i, t_str));
    fprintf(asmout, "PUSH\t%s\n", REGISTER_STRING[EAX]);
    fprintf(asmout, "PUSH\tOFFSET %s\n", "c_str_int_in");

    release_tmp_reg();
    fprintf(asmout, "CALL\tscanf\n");
    fprintf(asmout, "ADD\tESP, 0%xH\n", 8);
}
// From midcode to assembly code.
static void mid_to_asm(int mi)
{
    static int push_count = 0;
    int i = mi;
    switch (midcode_table[i].op0) {
        case ADD_OP:
            add_inst(i);
            break;
        case MINUS_OP:
            sub_inst(i);
            break;
        case MUL_OP:
            mul_inst(i);
            break;
        case DIV_OP:
            div_inst(i);
            break;
        case IFLT_OP:
        case IFLET_OP:;
        case IFMT_OP:
        case IFMET_OP:
        case IFEQ_OP:
        case IFNE_OP:
            if_inst(i, midcode_table[i].op0);
            break;
        case PARA_OP:
            // Count needed.
            push_inst(i);
            push_count++;
            break;
        case CALL_OP:
            call_inst(i);
            fprintf(asmout, "ADD\tESP, 0%xH\n", push_count << 2);
            push_count = 0;
            break;
        case JMP_OP:
            jmp_inst(i);
            break;
        case ASSIGN_OP:
            mov_inst(i);
            break;
        case PRINTF_OP:
            prf_inst(i);
            break;
        case SCANF_OP:
            scf_inst(i);
            break;
        case RETV_OP:
            rev_inst(i);
            break;
        case RETN_OP:
            ret_inst(i);
            break;
        default:
            break;
    }
}
// @param: func, func_mid_map
static void gen_func_code(int fi, int func[/*2*/])
{
    int i, j;
    int fp = fi + 1;
    fprintf(asmout, ";############### PROC %s ################\n", func_table[fi].name);
    // Function head.
    fprintf(asmout, "_%s\tPROC\n", midcode_table[func[0]].op1.operandCont.funcOp->name);
    // Initialize running stack.
    //push        ebp  
    //mov         ebp,esp 
    //sub         esp,0E4h 
    //push        ebx  
    //push        esi  
    //push        edi  
    //lea         edi,[ebp-0E4h] 
    //mov         ecx,39h 
    //mov         eax,0CCCCCCCCh 
    //rep stos    dword ptr es:[edi] 
    fprintf(asmout, "PUSH\tEBP\n");
    fprintf(asmout, "MOV\tEBP, ESP\n");
    fprintf(asmout, "SUB\tESP, 0%xH\n", func_var_num[fp] << 2);
    fprintf(asmout, "PUSH\tEBX\n");
    fprintf(asmout, "PUSH\tESI\n");
    fprintf(asmout, "PUSH\tEDI\n");

    i = 0;
    while (bblock[fi][i][0] !=0 || bblock[fi][i][1] != 0) {
        fprintf(asmout, ";BBLOCK_%d\n", i);
        fprintf(asmout, "TARGET_%d:\n", bblock[fi][i][0]);
        for (j = bblock[fi][i][0]; j <= bblock[fi][i][1]; j++) {
            mid_to_asm(j);
        }
        release_tmp_reg();
        i++;
    }
    fprintf(asmout, ";BBLOCK_%d\n", i);
    fprintf(asmout, "TARGET_%d:\n", bblock[fi][i - 1][1] + 1);
    fprintf(asmout, "TARGET_END:\n");
    // Clear running stack.
    //pop         edi  
    //pop         esi  
    //pop         ebx  
    //mov         esp,ebp 
    //pop         ebp  
    //ret              
    fprintf(asmout, "POP\tEDI\n");
    fprintf(asmout, "POP\tESI\n");
    fprintf(asmout, "POP\tEBX\n");
    fprintf(asmout, "MOV\tESP, EBP\n");
    fprintf(asmout, "POP\tEBP\n");
    fprintf(asmout, "RET\n");
    // Function end.
    fprintf(asmout, "_%s\tENDP\n", midcode_table[func[0]].op1.operandCont.funcOp->name);
}
static void gen_global_var()
{
    int i;
    print_line(";############# Global Variable ##############");
    for (i = 0; i < normal_table_p; i++) {
        if (normal_table[i].scope == NULL) {
            fprintf(asmout, "_%s\tDWORD\t%d\n", normal_table[i].name, 0);
        }
    }
}

static void gen_const_string()
{
    int i;
    print_line(";############# Const String #################");
    for (i = 0; i < midcode_table_p; i++) {
        if (midcode_table[i].op0 == PRINTF_OP && midcode_table[i].op1.operandType == STR_MB) {
            fprintf(asmout, "c_str_%d\tSBYTE\t\"%s\", 0dH, 0aH, 0\n", i, midcode_table[i].op1.operandCont.stringOp);
        }
    }
}
static void gen_head()
{
    print_line(";############# Assembly Code ################");
    print_line(".386");
    print_line(".MODEL FLAT");
    print_line("");
    print_line("OPTION casemap:none");
    print_line("includelib msvcrt.lib");
    print_line("printf proto C:dword,:vararg");
    print_line("scanf proto C:dword,:vararg");
    print_line("_ftol proto C:dword");
    print_line("");
    print_line(".STACK 4096");
    print_line("");
}

static void gen_data()
{
    fprintf(asmout, ".DATA\n");
    gen_global_var();
    gen_const_string();

    // For calling printf and scanf.
    fprintf(asmout, "%s\tSBYTE\t%s, 0dH, 0aH, 0\n", "c_str_char_out", "\"%c\"");
    fprintf(asmout, "%s\tSBYTE\t%s, 0dH, 0aH, 0\n", "c_str_int_out", "\"%d\"");
    fprintf(asmout, "%s\tSBYTE\t%s, 0dH, 0aH, 0\n", "c_str_float_out", "\"%f\"");

    fprintf(asmout, "%s\tSBYTE\t%s, 0\n", "c_str_char_in", "\"%c\"");
    fprintf(asmout, "%s\tSBYTE\t%s, 0\n", "c_str_int_in", "\"%d\"");
    fprintf(asmout, "%s\tSBYTE\t%s, 0\n", "c_str_float_in", "\"%f\"");
    print_line("");
}

static void gen_code()
{
	int i;
    fprintf(asmout, ".CODE\n");
	for (i = 0; i < func_table_p; i++) {
		init_reg_pool();
        global_reg_dist(func_mid_map[i]);
		gen_func_code(i, func_mid_map[i]);
	}
	fprintf(asmout, "\tEND\t_MAIN\n");
}

void gen_assembly_code()
{
	gen_bblock();
	gen_stack_ref();

    gen_head();
    gen_data();
    gen_code();
}